#include <opencv2/core/core.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/core/types_c.h>
#include <opencv2/highgui/highgui.hpp>
#include <iostream>
#include <fstream>
#include <cmath>

//#include <Windows.h>
#include <thread>

#include "mat_utils.h"
#include "GVFField.h"

using namespace cv;
using namespace std;

GVFField::GVFField(string in_file) {
	ifstream file;
	file.open(in_file, ios::in);
	file >> *(this);
}

GVFField::GVFField(string in_img, grad_method method, int in_iter, double in_mu, double in_factor)
{
	int scale = 1;
	int delta = 0;
	int ddepth = CV_64F;
	int iter = in_iter;
	double mu = in_mu;

	/// Load an image
	if (method == STORED_GVF) { in_img = "room.pgm"; }
	src = imread( in_img );
	imgData *id = new imgData("base image", src);
	std::thread dispThread(dispImg, id);
	dispThread.detach();


	//TODO raise error handle
	//if( !src.data )
	//{ return -1; }

	GaussianBlur( src, src, Size(3,3), 0, 0, BORDER_DEFAULT );

	/// Convert it to gray
	cvtColor( src, src_gray, CV_RGB2GRAY );

	/// Create window
	//namedWindow( window_name, CV_WINDOW_AUTOSIZE );

	/// Generate grad_x and grad_y
	Mat temp;
	Mat grad_x, grad_y;
	Mat abs_grad_x, abs_grad_y;
	
	Mat sq_grad;
	Mat f_x, f_y;
	//Mat abs_f_x, abs_f_y;
	Mat sq_f_x, sq_f_y;
	Mat b, c1, c2;
	//Mat u, v;

	/// Gradient X
	switch (method) { 
		case SCHARR_METHOD: 
			Scharr( src_gray, grad_x, ddepth, 1, 0, scale, delta, BORDER_REFLECT101 );
			Scharr( src_gray, grad_y, ddepth, 0, 1, scale, delta, BORDER_REFLECT101 );
			break;
		case SOBEL_METHOD: 
			Sobel( src_gray, grad_x, ddepth, 1, 0, 3, scale, delta, BORDER_DEFAULT );
			Sobel( src_gray, grad_y, ddepth, 0, 1, 3, scale, delta, BORDER_DEFAULT );
			break;
		default: 
			Sobel( src_gray, grad_x, ddepth, 1, 0, 3, scale, delta, BORDER_DEFAULT );
			Sobel( src_gray, grad_y, ddepth, 0, 1, 3, scale, delta, BORDER_DEFAULT );
			break;
	}
	//converts Mat back into 8bit unsigned
	convertScaleAbs( grad_x, abs_grad_x );
	convertScaleAbs( grad_y, abs_grad_y );
	
	/// Gradient Y

	/// Total Gradient (approximate)
	addWeighted( abs_grad_x, 0.5, abs_grad_y, 0.5, 0, grad );

	grad.convertTo(temp, ddepth, 1.0/255.0);
	grad = temp.clone();

	//logfile.setf(std::ios::scientific, std:: ios::floatfield);
	//logfile.precision(6);
	//for (int y = 0; y < temp.rows; ++y) {
	//  for (int x = 0;x < temp.cols; ++x) {
	//	  logfile<<temp.at<double>(y, x)<<" ";
	//  }
	//  logfile<<endl;
	//}
	
	imgData gradImgData = imgData("gradient", grad.clone());
	std::thread dispThread2(dispImg, &gradImgData);
	dispThread2.detach();

	//Square the Intensity gradient elements = Edge Map
	cv::pow(grad, 2, sq_grad);
	
	//Compute Partial Derivs of Edge Map
	//NOTE - Check Sobel kernel and del_t agreement (see NOTE on del_t)
	switch (method) { 
		case SCHARR_METHOD: 
			Scharr( sq_grad, f_x, ddepth, 1, 0, scale, delta, BORDER_REFLECT101 );
			Scharr( sq_grad, f_y, ddepth, 0, 1, scale, delta, BORDER_REFLECT101 );
			break;
		case SOBEL_METHOD: 
			Sobel( sq_grad, f_x, ddepth, 1, 0, 3, scale, delta, BORDER_REFLECT101 );
			Sobel( sq_grad, f_y, ddepth, 0, 1, 3, scale, delta, BORDER_REFLECT101 );
			break;
		default: 
			Sobel( sq_grad, f_x, ddepth, 1, 0, 3, scale, delta, BORDER_REFLECT101 );
			Sobel( sq_grad, f_y, ddepth, 0, 1, 3, scale, delta, BORDER_REFLECT101 );
			break;
	}
	
	//Square Partial Derivs
	cv::pow(f_x, 2, sq_f_x);
	cv::pow(f_y, 2, sq_f_y);
	//Compute iteration constants
	b = sq_f_x + sq_f_y;
	c1 = f_x.mul(b);
	c2 = f_y.mul(b);
	
	//Initialize the GVF Mats
	u = f_x.clone();
	v = f_y.clone();
	
	//set borders to 0, Neumann boundary condition
	u.row(0) = 0;
	v.row(0) = 0;
	u.col(0) = 0;
	v.col(0) = 0;
	u.row(u.size().height-1) = 0;
	v.row(v.size().height-1) = 0;
	u.col(u.size().width-1) = 0;
	v.col(v.size().width-1) = 0;


	//NOTE - tempermental variable, GVF converges only for a handful of values depending on method used
	//Sobel Kernel 1 => 1.0/(8.0 * mu)
	//Sobel Kernel 3 => 1.0/(64.0 * mu)
	//Scharr Kernel  => 1.0/(1000.0 * mu)
	double factor = 1000;
	if (in_factor > 0) {
		factor = in_factor;
	}
	else if (method == SCHARR_METHOD) factor = 1000;
	else if (method == SOBEL_METHOD) factor = 64;

	double del_t = 1.0/(factor * mu);
	double r = mu * del_t;
	Mat b_delt, c1_delt, c2_delt;
	b_delt = -b.mul(del_t);
	c1_delt = c1.mul(del_t);
	c2_delt = c2.mul(del_t);



	//Const mask to be applied to the iteration
	Mat mask = (Mat_<double>(3,3) <<	0, r,		0,
										r, 1-(4*r),	r,
										0, r,		0);

	for (int i =0; i<iter; i++) {
		Mat temp_u, temp_v;
		
		filter2D(u, temp_u, u.depth(), mask);
		temp_u = temp_u + b_delt.mul(u) + c1_delt;
		
		filter2D(v, temp_v, v.depth(), mask);
		temp_v = temp_v + b_delt.mul(v) + c2_delt;
		
		u = temp_u.clone();
		v = temp_v.clone();
	}
	
	//logfile<<"u"<<endl<<u<<endl;
	
	CvScalar foreground;
	CvScalar background;
	foreground.val[0] = 10;
	foreground.val[1] = 10;
	foreground.val[2] = 10;
	foreground.val[3] = 10;
	background.val[0] = 0;
	background.val[1] = 0;
	background.val[2] = 0;
	background.val[3] = 0;

	//TODO does not work, need to fix
	//draw_flow_image(u, v, src, background, foreground, 0.1);

	//reads in a stored precomputed GVF file
	if (method == STORED_GVF) {
		//ofstream fieldndxs;
		//fieldndxs.open("C:\\Users\\nsmith\\Dropbox\\1Research\\Snake Project\\Snake_Stepper\\Snake_Stepper\\fieldndxs.txt", ios::out);
		//TEMP read in GVF from file
		ifstream ufile, vfile;
		ufile.open("u2.txt", ios::in);
		vfile.open("v2.txt", ios::in);

		u = Mat(src.rows, src.cols, ddepth, Scalar(0.0));
		v = Mat(src.rows, src.cols, ddepth, Scalar(0.0));

		string uline, vline;
		int ndy = 0;
		while (getline(ufile, uline) && getline(vfile, vline)) {
			istringstream uiss(uline);
			istringstream viss(vline);
			double udbl, vdbl = 0;
			vector<double> uvec, vvec;
			int ndx = 0;
			while ((uiss >> udbl)&&(viss >> vdbl)) {
				u.at<double>(ndx,ndy) = udbl;
				v.at<double>(ndx,ndy) = vdbl;
				ndx++;
			}
			ndy++;
		}
	}

}


GVFField::~GVFField(void)
{
}


std::ostream &operator<<(std::ostream &pipe, GVFField &out) {
	int rows = out.u.rows;
	int cols = out.u.cols * out.u.channels();
	
	/* can't do this since we need the endl to determine end of row
	if (out.u.isContinuous()) {
		cols *= rows;
		rows = 1;
	}
	*/

	pipe << "U: " << rows << " " << out.u.cols << " " << out.u.channels() << endl;

	double *p;
	int i, j = 0;
	for (i = 0; i < rows; i++) {
		p = out.u.ptr<double>(i);
		for (j = 0; j < cols; j++) {
			pipe << p[j] << " ";
		}
		pipe << endl;
	}
	
	pipe << "V:" << endl;
	
	for (i = 0; i < rows; i++) {
		p = out.v.ptr<double>(i);
		for (j = 0; j < cols; j++) {
			pipe << p[j] << " ";
		}
		pipe << endl;
	}

	return pipe;
}


//NOTE: this is designed for single channel GVF only
std::istream &operator>>(std::istream &pipe, GVFField &in) {
	int ddepth = CV_64F;
	
	string check = "";
	//getline(pipe, check);
	pipe >> check;

	if (strcmp(check.c_str(), "U:") == 0) {
		int rows, cols, channels = 0;
		pipe >> rows >> cols >> channels;

		in.u = Mat(rows, cols, ddepth);
		in.v = Mat(rows, cols, ddepth);

		double *p;
		int i, j = 0;
		for (i = 0; i < rows; i++) {
			p = in.u.ptr<double>(i);
			for (j = 0; j < cols; j++){
				pipe >> p[j];
			}
		}

		pipe >> check;
		if (strcmp(check.c_str(), "V:") != 0) { throw new exception(); }

		for (i = 0; i < rows; i++) {
			p = in.v.ptr<double>(i);
			for (j = 0; j < cols; j++){
				pipe >> p[j];
			}
		}
	}
	else {
		//did not contain the expected entry
		throw new exception();
	}

	return pipe;
}

void GVFField::normalize(double min, double max) {
	double range = max - min;

	double umax= 0.0;
	double umin = 0.0;
	for (MatIterator_<double> it = u.begin<double>(); it != u.end<double>(); it++) {
		if ((*it) > umax) {
			umax = (*it);
		}
		else if ((*it) < umin) {
			umin = (*it);
		}
	}
	double urange = umax - umin;
	for (MatIterator_<double> it = u.begin<double>(); it != u.end<double>(); it++) {
		(*it) = ((*it) - umin)/urange;
		(*it) = ((*it)*range+min);
	}

	double vmax= 0.0;
	double vmin = 0.0;
	for (MatIterator_<double> it = v.begin<double>(); it != v.end<double>(); it++) {
		if ((*it) > vmax) {
			vmax = (*it);
		}
		else if ((*it) < vmin) {
			vmin = (*it);
		}
	}
	double vrange = vmax - vmin;
	for (MatIterator_<double> it = v.begin<double>(); it != v.end<double>(); it++) {
		(*it) = ((*it) - vmin)/vrange;
		(*it) = ((*it)*range+min);
	}
}